package org.king.util;

import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.stereotype.Component;

import java.util.Optional;
import java.util.concurrent.TimeUnit;

/**
 * 限流通用工具
 * 针对于不同的情况进行限流操作，支持限流升级
 */
@Slf4j
@Component
public class FlowUtil {

    @Resource
    private StringRedisTemplate template;

    public boolean limitOnceCheck(String key, int blockTime){
        if(Boolean.TRUE.equals(template.hasKey(key))){
            return false;
        }else{
            template.opsForValue().set(key, "", blockTime, TimeUnit.SECONDS);
        }
        return true;
    }

    /**
     * 针对于在时间段内多次请求限制，如3秒内限制请求20次，超出频率则封禁一段时间
     * @param counterKey 计数键
     * @param blockKey 封禁键
     * @param blockTime 封禁时间
     * @param frequency 请求频率
     * @param period 计数周期
     * @return 是否通过限流检查
     */
    public boolean limitPeriodCheck(String counterKey, String blockKey, int blockTime, int frequency, int period){
        return this.internalCheck(counterKey, frequency, period, (overclock) -> {
            if (overclock)
                template.opsForValue().set(blockKey, "", blockTime, TimeUnit.SECONDS);
            return !overclock;
        });
    }
    /**
     * 内部使用请求限制主要逻辑
     * @param key 计数键
     * @param frequency 请求频率
     * @param period 计数周期
     * @param action 限制行为与策略
     * @return 是否通过限流检查
     */
    private boolean  internalCheck(String key, int frequency, int period, LimitAction action) {
        // 1. 获取当前计数
        String count = template.opsForValue().get(key);

        if (count != null) {
            // 2. 计数存在时，递增并检查是否超频
            long value = Optional.ofNullable(template.opsForValue().increment(key)).orElse(0L);
            int c = Integer.parseInt(count);

            // 3. 修复可能的并发问题：确保递增后刷新过期时间
            if (value != c + 1)
                template.expire(key, period, TimeUnit.SECONDS);

            // 4. 执行限流策略（通过 Lambda 判断是否超频）
            return action.run(value > frequency);
        } else {
            // 5. 首次请求时初始化计数
            template.opsForValue().set(key, "1", period, TimeUnit.SECONDS);
            return true; // 允许通过
        }
    }
    /**
     * 内部使用，限制行为与策略
     */
    private interface LimitAction {
        boolean run(boolean overclock);
    }
}
